from jaxtyping import Array, Float, Int, Bool, PyTree
import jax
import jax.numpy as jnp


def lnm_logp(x: Float[Array, ""],
            locs: Float[Array, "num_components"], 
            log_scales: Float[Array, "num_components"],
            log_weights: Float[Array, "num_components"]):
    logx = jnp.log(x+1e-8)
    temp = jax.scipy.stats.norm.logpdf(logx, locs, jnp.exp(log_scales))
    temp = temp + log_weights
    temp = jax.nn.logsumexp(temp)
    ll = temp - logx
    return ll

def lnm_sample(locs: Float[Array, "num_components"], 
               log_scales: Float[Array, "num_components"],
               log_weights: Float[Array, "num_components"],
               key: Array, 
               num_samples: int):
    key1, key2= jax.random.split(key)
    eps = jax.random.normal(key1, (num_samples, ))
    clusters = jax.random.categorical(key2, log_weights, shape=(num_samples, ))
    locs = locs[clusters] # (num_samples, )
    log_scales = log_scales[clusters] # (num_samples, )
    samples = jnp.exp(eps * jnp.exp(log_scales) + locs)
    return samples # (num_samples, )
